import math
import random
import numpy as np
import argparse

class PSO:
        #对象内置构造方法，初始化时自定义
    def __init__(self, dimension, time, size, low, up, v_low, v_high):
        # 初始化
        self.dimension = dimension  # 变量个数  粒子
        self.time = time  # 迭代的代数
        self.size = size  # 种群大小    
        self.bound = []  # 变量的约束范围   
        self.bound.append(low)
        self.bound.append(up)
        self.v_low = v_low      #速度
        self.v_high = v_high
        #np.zeros创建一个用于存放粒子最优解以及群体最优解的数组
        self.x = np.zeros((self.size, self.dimension))  # 所有粒子的位置
        self.v = np.zeros((self.size, self.dimension))  # 所有粒子的速度
        #pbest 个体最佳位置 gbest群体最佳位置
        self.p_best = np.zeros((self.size, self.dimension))  # 每个粒子最优的位置
        self.g_best = np.zeros((1, self.dimension))[0]  # 全局最优的位置

        # 初始化第0代初始全局最优解
        temp = float('inf')
        for i in range(self.size):
            for j in range(self.dimension):
                self.x[i][j] = random.uniform(self.bound[0][j], self.bound[1][j])
                self.v[i][j] = random.uniform(self.v_low, self.v_high)
            self.p_best[i] = self.x[i]  # 储存最优的个体
            fit = self.fitness(self.p_best[i])
            # 做出修改
            if abs(fit) > temp:
                self.g_best = self.p_best[i]
                temp = fit

    def fitness(self, x):
        """
        函数计算
        """
        #函数计算的公式
        #这个不影响答案结果 影响结果的因素为粒子群算法的初始定义的几个值
        x0 = x[0]
        y = math.cos(x0) - x0
        return y

    def update(self, size): #更新粒子速度和位置
        c1 = 2.0  # 粒子个体学习因子
        c2 = 2.0  #粒子的社会学习因子
        w = 0.9  # 自身权重因子 （速度惯性权重）
        for i in range(size):
            # 更新速度(核心公式)
            # 即为这只鸟第i步的速度为=上一步自身的速度惯性+自我认知部分+社会认知部分
            #v_{i}^{d}=wv_{i}^{d-1}+c_1r_1(pbest_{i}^{d}-x_{i}^{d})+c_2r_2(gbest^d-x_{i}^{d})\\
            self.v[i] = w * self.v[i] + c1 * random.uniform(0, 3.14) * (
                    self.p_best[i] - self.x[i]) + c2 * random.uniform(0, 3.14) * (self.g_best - self.x[i])
            # 速度限制
            for j in range(self.dimension):
                #取速度值               
                if self.v[i][j] < self.v_low:
                    self.v[i][j] = self.v_low
                if self.v[i][j] > self.v_high:
                    self.v[i][j] = self.v_high
            # 更新位置
            self.x[i] = self.x[i] + self.v[i]
            # 位置限制
            for j in range(self.dimension):
                if self.x[i][j] < self.bound[0][j]:
                    self.x[i][j] = self.bound[0][j]
                if self.x[i][j] > self.bound[1][j]:
                    self.x[i][j] = self.bound[1][j]
            # 更新p_best和g_best
            if abs(self.fitness(self.x[i])) < abs(self.fitness(self.p_best[i])):
                self.p_best[i] = self.x[i]
            if abs(self.fitness(self.x[i])) < abs(self.fitness(self.g_best)):
                self.g_best = self.x[i]

    def pso(self):
        self.final_best = np.array([random.uniform(0, 3.14)])
        for _ in range(self.time):
            self.update(self.size)
            if abs(self.fitness(self.g_best)) < abs(self.fitness(self.final_best)):
                self.final_best = self.g_best.copy()
            temp = self.fitness(self.final_best)
        print('当前最佳位置(x)：', self.final_best[0])
        print('当前的最佳适应度(y)：', temp)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-t', dest='times')
    parser.add_argument('-s', dest='size')
    args = parser.parse_args()
    # if not (args.times):
    #     parser.print_help()
    #     exit()
    time = 1
    size = 50
    dimension = 1
    v_low = -0.001
    v_high = 0.005
    low = [0.01]
    up = [3.14]
    pso = PSO(dimension, time, size, low, up, v_low, v_high)
    pso.pso()

